import asyncio
import os
import time
import concurrent.futures
from pydub import AudioSegment
import edge_tts

VOICE = "zh-CN-YunxiNeural"
TEMP_DIR = "temp_audio"
MAX_WORKERS = 4  # 同时生成的最大音频数

def parse_srt(file_path):
    """Parse SRT file and return list of (text, start_ms, end_ms) tuples"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    segments = []
    for segment in content.strip().split('\n\n'):
        lines = segment.split('\n')
        if len(lines) < 3:
            continue
            
        timecodes = lines[1].split(' --> ')
        if len(timecodes) != 2:
            continue
            
        def time_to_ms(time_str):
            hh_mm_ss, ms = time_str.split(',')
            h, m, s = hh_mm_ss.split(':')
            return (int(h)*3600 + int(m)*60 + int(s))*1000 + int(ms)
            
        start_ms = time_to_ms(timecodes[0])
        end_ms = time_to_ms(timecodes[1])
        text = '\n'.join(lines[2:]).strip()
        segments.append((text, start_ms, end_ms))
    
    return segments

def generate_segment_sync(text, temp_file):
    """Generate audio segment with retry on failure"""
    max_retries = 3
    retry_delay = 1
    
    for attempt in range(max_retries):
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            communicate = edge_tts.Communicate(text, VOICE,rate ="+45%")
            loop.run_until_complete(communicate.save(temp_file))
            return AudioSegment.from_mp3(temp_file)
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            print(f"生成失败 (尝试 {attempt + 1}/{max_retries}), 等待重试...")
            time.sleep(retry_delay)
        finally:
            loop.close()

async def process_srt(srt_file, output_file):
    """Process SRT with precise timing and parallel generation"""
    start_time = time.time()
    os.makedirs(TEMP_DIR, exist_ok=True)
    segments = parse_srt(srt_file)
    combined = AudioSegment.silent(duration=0)
    last_end = 0
    total = len(segments)
    thread_times = []
    
    # Generate all audio segments in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = []
        for i, (text, start_ms, end_ms) in enumerate(segments):
            temp_file = os.path.join(TEMP_DIR, f"temp_{i}.mp3")
            print(f"提交第 {i+1}/{total} 段音频生成任务...")
            futures.append(executor.submit(generate_segment_sync, text, temp_file))
        
        # Wait for all tasks to complete and process results
        for i, (future, (text, start_ms, end_ms)) in enumerate(zip(futures, segments)):
            segment_start = time.time()
            print(f"等待第 {i+1}/{total} 段音频完成...")
            audio = future.result()
            duration = time.time() - segment_start
            thread_times.append(duration)
            print(f"第 {i+1}/{total} 段音频完成! 耗时: {duration:.3f}秒")
            
            if start_ms > last_end:
                combined += AudioSegment.silent(duration=start_ms - last_end)
            combined += audio
            last_end = start_ms + len(audio)
    
    combined.export(output_file, format="mp3")
    total_time = time.time() - start_time
    thread_total = sum(thread_times)
    print(f"Generated synchronized audio saved to {output_file}")
    print(f"程序运行时间: {total_time:.2f}秒")
    print(f"线程总耗时: {thread_total:.2f}秒 (并行效率: {thread_total/total_time:.1f}x)")

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 3:
        print("Usage: python main.py input.srt output.mp3")
        sys.exit(1)
        
    srt_file = sys.argv[1]
    output_file = sys.argv[2]
    
    try:
        asyncio.run(process_srt(srt_file, output_file))
    except Exception as e:
        print(f"Error: {e}")
